def load_transactions():
    """加载示例交易数据"""
    return [
        ['牛奶', '面包', '啤酒'],
        ['牛奶', '面包', '咖啡'],
        ['牛奶', '尿布'],
        ['面包', '啤酒'],
        ['面包', '咖啡', '尿布'],
        ['牛奶', '面包', '啤酒', '尿布'],
        ['面包', '尿布']
    ]

def create_c1(transactions):
    """创建1-项集"""
    c1 = []
    for transaction in transactions:
        for item in transaction:
            if [item] not in c1:
                c1.append([item])
    c1.sort()
    return list(map(frozenset, c1))

def scan_dataset(dataset, candidates, min_support):
    """
    扫描数据集，计算支持度并筛选出满足最小支持度的项集
    返回：满足条件的项集和它们的支持度
    """
    ss_cnt = {}
    for tid in dataset:
        for can in candidates:
            if can.issubset(tid):
                if can not in ss_cnt:
                    ss_cnt[can] = 1
                else:
                    ss_cnt[can] += 1
    
    num_items = float(len(dataset))
    ret_list = []
    support_data = {}
    
    for key in ss_cnt:
        support = ss_cnt[key] / num_items
        if support >= min_support:
            ret_list.insert(0, key)
        support_data[key] = support
    
    return ret_list, support_data

def apriori_gen(freq_sets, k):
    """由k-1项集生成k项集"""
    ret_list = []
    len_lk = len(freq_sets)
    
    for i in range(len_lk):
        for j in range(i + 1, len_lk):
            # 前k-2个元素相同的项集合并
            l1 = list(freq_sets[i])[:k-2]
            l2 = list(freq_sets[j])[:k-2]
            l1.sort()
            l2.sort()
            
            if l1 == l2:
                # 合并两个项集
                ret_list.append(freq_sets[i] | freq_sets[j])
    
    return ret_list

def apriori(transactions, min_support=0.5):
    """
    Apriori算法主函数
    transactions: 交易数据集
    min_support: 最小支持度阈值
    返回：所有频繁项集和它们的支持度
    """
    # 创建1-项集并筛选
    c1 = create_c1(transactions)
    dataset = list(map(set, transactions))
    l1, support_data = scan_dataset(dataset, c1, min_support)
    
    # 存储所有频繁项集
    l = [l1]
    k = 2
    
    # 迭代生成更高阶的频繁项集
    while len(l[k-2]) > 0:
        # 生成候选k-项集
        ck = apriori_gen(l[k-2], k)
        # 筛选出满足支持度的k-项集
        lk, sup_k = scan_dataset(dataset, ck, min_support)
        
        support_data.update(sup_k)
        l.append(lk)
        k += 1
    
    return l, support_data

def generate_rules(l, support_data, min_confidence=0.7):
    """
    从频繁项集中生成关联规则
    l: 所有频繁项集
    support_data: 项集的支持度
    min_confidence: 最小置信度阈值
    返回：满足条件的关联规则
    """
    rules = []
    # 只处理长度大于1的频繁项集
    for i in range(1, len(l)):
        for freq_set in l[i]:
            # 生成所有非空子集
            h1 = [frozenset([item]) for item in freq_set]
            
            if i > 1:
                # 对于长度大于2的项集，使用更复杂的方法生成规则
                rules_from_conseq(freq_set, h1, support_data, rules, min_confidence)
            else:
                # 计算置信度并筛选
                calc_confidence(freq_set, h1, support_data, rules, min_confidence)
    
    return rules

def calc_confidence(freq_set, h, support_data, rules, min_confidence):
    """计算置信度并筛选规则"""
    pruned_h = []
    for conseq in h:
        # 置信度 = 支持度(freq_set) / 支持度(freq_set - conseq)
        conf = support_data[freq_set] / support_data[freq_set - conseq]
        if conf >= min_confidence:
            print(freq_set - conseq, "=>", conseq, "置信度:", conf)
            rules.append((freq_set - conseq, conseq, conf))
            pruned_h.append(conseq)
    return pruned_h

def rules_from_conseq(freq_set, h, support_data, rules, min_confidence):
    """从更长的项集中生成规则"""
    m = len(h[0])
    if len(freq_set) > (m + 1):
        # 生成m+1长度的候选规则
        hmp1 = apriori_gen(h, m + 1)
        hmp1 = calc_confidence(freq_set, hmp1, support_data, rules, min_confidence)
        
        if len(hmp1) > 1:
            # 递归生成更长的规则
            rules_from_conseq(freq_set, hmp1, support_data, rules, min_confidence)

# 示例用法
if __name__ == "__main__":
    # 加载交易数据
    transactions = load_transactions()
    
    # 应用Apriori算法，最小支持度设为0.3
    l, support_data = apriori(transactions, min_support=0.3)
    print("\n频繁项集:")
    for i, item_set in enumerate(l):
        if item_set:  # 只打印非空项集
            print(f"{i+1}-项集: {item_set}")
    
    print("\n支持度:")
    for item, support in support_data.items():
        if len(item) > 0:
            print(f"{item}: {support:.4f}")
    
    # 生成关联规则，最小置信度设为0.6
    print("\n关联规则:")
    rules = generate_rules(l, support_data, min_confidence=0.6)
